package bayes;

import java.util.Iterator;
import java.util.Vector;
import java.util.Map.Entry;
import bayes.nodeTypes.Variable;
import domain.Debug;

public class Enumeration {

	public static Distribution runEnumerationAlgorithm(Variable x, Evidence evidence, BayesNetwork bn) {

		Distribution ans =  enumeration_ask(x, evidence, bn);
		System.out.println("\nP(" + x.get_name() + " | " + allEvidencesToString(evidence) + ") = " + ans);
		return ans;
	}



	private static String allEvidencesToString(Evidence evidence) {
		String retStr = "";

		Iterator<Entry<Variable, Integer>> iterEvidence = evidence.iterator();
		while(iterEvidence.hasNext()) {
			Entry<Variable, Integer> varNval = iterEvidence.next();
			Variable var = varNval.getKey();

			// variable name:
			String varName = var.get_name();
			
			// variable value:
			Integer intValue = varNval.getValue();
			String varValue = var.getDomainValName(intValue);
			
			retStr += varName+"="+varValue;
			if (iterEvidence.hasNext()) retStr += ", ";
		}
		return retStr;
	}



//		function Enumeration-Ask(X, e, bn) returns a distribution over X
//		inputs: X, the query variable
//			e, observed values for variables E
//			bn, a Bayes net with variables {X} u E u Y /* Y = hidden variables */
//
//		Q(X) <- a distribution over X, initially empty
//		for each value xi of X do
//			extend e with value xi for X
//			Q(xi) <- Enumerate-All(Vars[bn], e)
//		return Normalize(Q(X))
//
//	function Enumerate-All(vars, e) returns a real number
//		if Empty?(vars) then return 1.0
//		Y <- First(vars)
//		if Y has value y in e
//			then return P(y | parents(Y)) * Enumerate-All(Rest(vars), e)
//			else return the sum over y of P(y | parents(Y)) * Enumerate-All(Rest(vars), ey)
//				where ey is e extended with Y = y


	private static Distribution enumeration_ask(Variable variable, Evidence evidence, BayesNetwork bn){
		Distribution distribution = new Distribution(variable.get_domainSize());
		double enumerate_all;
		for (int i = 0; i < variable.get_domainSize(); i++) {
			Evidence newEvidence = new Evidence(evidence);
			newEvidence.extend(variable, i);
			Vector<Node> allVariables = bn.getAllVariables();
			Debug.println(allVariables);

			enumerate_all = enumerate_all(allVariables, newEvidence);

			distribution.setElementAt(i, enumerate_all);
		}
		distribution.normalize();
		return distribution;
	}

	private static double enumerate_all(Vector<Node> variablesNodes, Evidence evidence) {

		if (variablesNodes.isEmpty()){
			//Debug.println("ans = 1.0");
			return 1.0;
		}

		Node Y = variablesNodes.remove(0);
		Vector<Node> restVars = variablesNodes;
		//Debug.println("Y: "+Y);
		//Debug.println("Rest: "+restVars);
		//Debug.println(evidence);

		double ans = 0;

		double yGivenParents;
		int parentDomainValue;
		if (evidence.contains(Y.getVariable())){

			// ENUMERATE-ALL(REST(vars), e):
			double nextEnumAll = enumerate_all(restVars, evidence);

			// P(y|Parents(Y)):
			int y = evidence.getVarValue(Y.getVariable());
			parentDomainValue = getParentValue(evidence, Y);
			yGivenParents = Y.getProbability(parentDomainValue, y);

			// ans = P(y|Parents(Y)) * ENUMERATE-ALL(REST(vars), e):
			ans = yGivenParents * nextEnumAll;
			//Debug.println("ans = " + ans);
			return ans;
		} else {
			ans = 0;
			parentDomainValue = getParentValue(evidence, Y);
			for (int domainVal = 0; domainVal < Y.get_domainSize(); domainVal++) {

				// P(y|Parents(Y)):
				yGivenParents = Y.getProbability(parentDomainValue, domainVal);

				// ENUMERATE-ALL(REST(vars), e_y):
				Evidence evidence_extended = new Evidence(evidence);
				evidence_extended.extend(Y.getVariable(), domainVal);
				Vector<Node> restVarsClone = new Vector<Node>(restVars);
				double nextEnumAll = enumerate_all(restVarsClone, evidence_extended);

				// ans = ans + ( P(y|Parents(Y)) * ENUMERATE-ALL(REST(vars), e_y) ):
				ans += yGivenParents * nextEnumAll;
			}
			//Debug.println("ans = " + ans);
			return ans;
		}
	}

	private static int getParentValue(Evidence evidence, Node Y) {
		int parentValue;
		if (Y.hasParents()){
			parentValue = evidence.getVarValue(Y.get_parent().getVariable());
		}else{
			parentValue = 0;
		}
		return parentValue;
	}
}